#coding=utf-8
import cv2
import sys
import os

cap = cv2.VideoCapture(0)
# while 1:
#     _ , framee = cap.read()
#     cv2.imshow("test",framee)
#     cv2.waitKey(10)
# cv2.destroyAllWindows
# cap.release
class_count = 0
flag=1
while flag:
    class_name=raw_input("请输入标签类别名称：")
    print("标签类别为："+class_name+"通过空格键采集图片")
    os.mkdir("train_dir/"+class_name)
    class_dir="train_dir/"+class_name+"/"
    count=0
    while True:
        # 从摄像头读取图片
        success, img = cap.read()
        # img = cv2.flip(img, 1)
        img2=img.copy()
        img = cv2.rectangle(img, (212, 112), (468, 368), (0, 255, 0), 2)
        data_img = img2[112:368,212:468]

        try:
            cv2.imshow('image', img)   # imshow文件不能是中文，处理报错
        except Exception as e:
            pass

        k = cv2.waitKey(1)
        if k == ord(' '):
            cv2.imencode('.jpg', data_img)[1].tofile(class_dir+class_name+str(count)+'.jpg')
            count += 1

        if k == ord('q'):
            cap.release()
            cv2.destroyAllWindows()
            flag=0

        if k == ord('n'):
            break

